use std::{
    collections::HashSet,
    fs::{self, File},
    io::{Cursor, ErrorKind, Read, Seek},
    path::{Path, PathBuf},
    str::FromStr,
};

use anyhow::{Context, Result, bail};
use byteorder::{LittleEndian, ReadBytesExt};
use encoding_rs::GBK;
use hdf5::{File as H5File, Group, H5Type, types::VarLenUnicode};
use image::{GrayImage, Luma};
use zip::ZipArchive;

use crate::data::utils;

/// 把 labels 写入 txt，等价于 Python 的 write_txt
fn write_txt<P: AsRef<Path>>(path: P, lines: &[String]) -> Result<()> {
    utils::write_txt(path, lines, "\n")
}

/// 解析单个 gnt 文件，返回 Vec<(png_file_name, label)>
fn parse_gnt_file<R: Read + Seek, P: AsRef<Path>>(
    f: &mut R,
    images_dir: P,
    hdf5_grp: Option<Group>,
) -> Result<Vec<(String, String, u16, u16)>> {
    let mut results = Vec::new();
    let mut count = 0usize;

    loop {
        // 先读 4 字节长度
        let length = match f.read_u32::<LittleEndian>() {
            Ok(l) => l,
            Err(e) if e.kind() == ErrorKind::UnexpectedEof => break,
            Err(e) => bail!("read length failed: {}", e),
        };

        // 剩余长度 = length - 4(长度字段本身)
        let mut buf = vec![0u8; (length - 4) as usize];
        f.read_exact(&mut buf)?;

        let mut cur = Cursor::new(&buf);

        // tag_code: 2 字节
        let mut tag_bytes = [0u8; 2];
        cur.read_exact(&mut tag_bytes)?;
        let (label, _, _) = GBK.decode(&tag_bytes);
        let label = label.trim_end_matches('\0').to_string();

        let width = cur.read_u16::<LittleEndian>()?;
        let height = cur.read_u16::<LittleEndian>()?;

        if width == 0 || height == 0 {
            eprintln!("invalid image size: ({}, {})", width, height);
            continue;
        }

        // 读取 height * width 个像素
        let mut img = GrayImage::new(width as u32, height as u32);
        let mut img_data = Vec::new();
        for y in 0..height {
            for x in 0..width {
                let gray = cur.read_u8()?;
                img_data.push(gray);
                img.put_pixel(x as u32, y as u32, Luma([gray]));
            }
        }

        count += 1;

        let png_name = format!("{}.png", count);
        let raw_png_name: &str = &png_name;

        if let Some(ref hdf5_grp) = hdf5_grp {
            hdf5_grp
                .new_dataset::<u8>()
                .shape([img_data.len()])
                .create(raw_png_name)?
                .write(&img_data)?
        } else {
            // 立即保存图片
            let out_dir = PathBuf::from(images_dir.as_ref());
            fs::create_dir_all(&out_dir)?;
            img.save(out_dir.join(&png_name))?;
        }

        results.push((png_name, label, width, height));
    }

    Ok(results)
}

pub fn batch_transfer(raw_dir: &str, images_dir: &str) -> Result<()> {
    let save_dir = PathBuf::from(images_dir);
    let mut label_list = Vec::new();
    for entry in fs::read_dir(raw_dir)? {
        let dir_entry = entry?;
        let gnt_path = dir_entry.path();
        if gnt_path.extension().map_or(true, |ext| ext != "gnt") {
            continue;
        }

        println!("Transfering {}...", gnt_path.display());
        let mut f = File::open(&gnt_path)?;
        let save_dir = save_dir.join(gnt_path.file_stem().context("invalid filename")?);
        let records = parse_gnt_file(&mut f, &save_dir, None)?;
        for (png, label, _, _) in records {
            let stem = gnt_path.file_stem().unwrap();
            let rel_path = PathBuf::from(stem).join(&png);
            label_list.push(format!("{}\t{}", rel_path.display(), label));
        }
    }

    write_txt(save_dir.join("labels.txt"), &label_list)?;
    println!("done");

    Ok(())
}

#[derive(H5Type, Clone, PartialEq, Debug)] // register with HDF5
#[repr(C)]
pub struct ImageMeta {
    pub name: VarLenUnicode,
    pub label: VarLenUnicode,
    pub width: u16,
    pub height: u16,
}

impl ImageMeta {
    pub fn new(name: &str, label: &str, width: u16, height: u16) -> Self {
        Self {
            name: VarLenUnicode::from_str(name).expect("can't convert name to varlen"),
            label: VarLenUnicode::from_str(label).expect("can't convert label to varlen"),
            width,
            height,
        }
    }
}

pub fn batch_unzip_and_transfer(
    zip_dir: &str,
    save_dir: &str,
    smart_check: bool,
    trn_list: Vec<String>,
    tst_list: Vec<String>,
    vld_list: Vec<String>,
    trn_tag: String,
    tst_tag: String,
    vld_tag: String,
    hdf5_filename: Option<&str>,
) -> Result<()> {
    let save_path = PathBuf::from(save_dir);
    let mut all_labels = HashSet::new();
    let hdf5_file = hdf5_filename.map(|filename| {
        H5File::create(save_path.join(filename)).expect("Failed to create HDF5 file")
    });
    let mut hdf5_trn_meta_list = Vec::new();
    let mut hdf5_tst_meta_list = Vec::new();
    let mut hdf5_vld_meta_list = Vec::new();
    let mut hdf5_ukn_meta_list = Vec::new();

    for entry in fs::read_dir(zip_dir)? {
        let mut label_list = Vec::new();
        let dir_entry = entry?;
        let zip_path = dir_entry.path();
        if zip_path.extension().map_or(true, |ext| ext != "zip") {
            continue;
        }

        println!("Unzip {}...", zip_path.display());
        // 1. 读取zip文件
        let file = File::open(&zip_path)?;
        let mut zip = ZipArchive::new(file)?;

        // 2. 准备输出根目录
        let zip_file_name = Path::new(zip_path.file_name().unwrap())
            .file_stem()
            .context("invalid filename")?
            .to_string_lossy()
            .to_string();

        let tag = if smart_check {
            let zip_file_name = zip_file_name.to_lowercase();
            if zip_file_name.contains(&trn_tag) {
                &trn_tag
            } else if zip_file_name.contains(&tst_tag) {
                &tst_tag
            } else if zip_file_name.contains(&vld_tag) {
                &vld_tag
            } else {
                "unknown"
            }
        } else {
            assert!(
                trn_list.contains(&zip_file_name)
                    || tst_list.contains(&zip_file_name)
                    || vld_list.contains(&zip_file_name),
                "For non-smart check mode, specify the compressed package that needs to be decompressed {}",
                zip_file_name
            );

            if trn_list.contains(&zip_file_name) {
                &trn_tag
            } else if tst_list.contains(&zip_file_name) {
                &tst_tag
            } else if vld_list.contains(&zip_file_name) {
                &vld_tag
            } else {
                "unknown"
            }
        };

        let tag_grp = hdf5_file.clone().map(|hdf5_file| {
            hdf5_file
                .group(tag)
                .or_else(|_| hdf5_file.create_group(tag))
                .expect("Failed to create group")
        });
        let tag_path = save_path.join(tag);

        let hdf5_subgrp = tag_grp.map(|tag_grp| {
            tag_grp
                .create_group(&zip_file_name)
                .expect("Failed to create group")
        });
        let root = PathBuf::from(tag_path.join(zip_file_name.clone()));

        if hdf5_filename.is_none() {
            std::fs::create_dir_all(&root)?;
        }

        for i in 0..zip.len() {
            let mut file = zip.by_index(i)?;
            let name = file.name().to_string();
            if !name.to_lowercase().ends_with(".gnt") {
                continue;
            }

            // 子目录名：去掉 .gnt 后缀
            let stem = Path::new(&name)
                .file_stem()
                .context("invalid filename")?
                .to_string_lossy()
                .to_string();
            let out_sub = root.join(&stem);
            let hdf5_out = hdf5_subgrp.clone().map(|hdf5_subgrp| {
                hdf5_subgrp
                    .create_group(&stem)
                    .expect("create hdf5 group failed")
            });

            if hdf5_filename.is_none() {
                std::fs::create_dir_all(&out_sub)?;
            }

            // 解析 gnt 并保存 png
            println!("Transfering {}...", &name);

            // 把 ZipFile 完整读到内存
            let mut buf = Vec::new();
            file.read_to_end(&mut buf)?;

            // 用 Cursor 包装，Cursor 实现了 Read + Seek
            let mut cursor = std::io::Cursor::new(buf);
            // 原有逻辑：保存为文件
            let recs = parse_gnt_file(&mut cursor, &out_sub, hdf5_out)?;
            for (png, label, width, height) in recs {
                let rel_path = PathBuf::from(&stem).join(&png);
                label_list.push(format!("{}\t{}", rel_path.display(), label));

                let meta_info = ImageMeta::new(
                    &format!("{}/{}/{}/{}", tag, zip_file_name, stem, png),
                    &label,
                    width,
                    height,
                );
                if tag == trn_tag {
                    hdf5_trn_meta_list.push(meta_info);
                } else if tag == tst_tag {
                    hdf5_tst_meta_list.push(meta_info);
                } else if tag == vld_tag {
                    hdf5_vld_meta_list.push(meta_info);
                } else if tag == "unknown" {
                    hdf5_ukn_meta_list.push(meta_info);
                } else {
                    panic!("unknown tag {}", tag);
                }

                all_labels.insert(label);
            }
        }

        if hdf5_filename.is_none() {
            // 只有在不保存到HDF5时才写入labels.txt文件
            write_txt(root.join("labels.txt"), &label_list)?;
        }
        println!("{} done", zip_path.display());
    }

    if let Some(hdf5_file) = hdf5_file {
        // 保存到HDF5文件
        let all_labels_ = all_labels
            .iter()
            .map(|label| {
                VarLenUnicode::from_str(&label).expect("Failed to convert label to VarLenUnicode")
            })
            .collect::<Vec<_>>();

        hdf5_file
            .new_dataset_builder()
            .with_data(&all_labels_)
            .create("classes")
            .expect("Failed to create dataset for classes");
        hdf5_file
            .new_dataset_builder()
            .with_data(&hdf5_trn_meta_list)
            .create("trn_meta")
            .expect("create hdf5 meta dataset failed");
        hdf5_file
            .new_dataset_builder()
            .with_data(&hdf5_tst_meta_list)
            .create("tst_meta")
            .expect("create hdf5 meta dataset failed");
        hdf5_file
            .new_dataset_builder()
            .with_data(&hdf5_vld_meta_list)
            .create("vld_meta")
            .expect("create hdf5 meta dataset failed");
        hdf5_file
            .new_dataset_builder()
            .with_data(&hdf5_ukn_meta_list)
            .create("ukn_meta")
            .expect("create hdf5 meta dataset failed");
    } else {
        // 原有逻辑：写入classes.txt文件
        write_txt(
            save_path.join("classes.txt"),
            &all_labels.into_iter().collect::<Vec<_>>(),
        )?;
    }
    Ok(())
}

#[cfg(test)]
pub mod tests {
    use super::*;
    #[test]
    fn test_hdf5() {
        let save_dir = r"D:\projects\datasets";
        let hdf5_filename = Some("hwdb.h5");
        let save_path = PathBuf::from(save_dir);
        let hdf5_file = hdf5_filename.map(|filename| {
            H5File::create(save_path.join(filename)).expect("Failed to create HDF5 file")
        });

        let mut hdf5_meta_list = Vec::new();
        for i in 0..10 {
            hdf5_meta_list.push(ImageMeta {
                name: VarLenUnicode::from_str(&format!("test {}", i))
                    .expect("Failed to create VarLenUnicode"),
                label: VarLenUnicode::from_str(&format!("label {}", i))
                    .expect("Failed to create VarLenUnicode"),
                width: i,
                height: i,
            });
        }

        if let Some(hdf5_file) = hdf5_file {
            hdf5_file
                .new_dataset_builder()
                .with_data(&hdf5_meta_list)
                .create("meta")
                .expect("create hdf5 meta dataset failed");
        }
    }
}
